from model import UNet
import torch
import torch.nn as nn
from dataset import BacteriaDataset, DataAug
from torch.utils.data import DataLoader
from torch.utils.data.sampler import RandomSampler
from tqdm import tqdm
from loss import DiceLoss, IOULoss

device = ("cuda" if torch.cuda.is_available() else "cpu")

num_epochs = 10
learning_rate = 0.001
batch_size = 16
shuffle = True
pin_memory = True
num_workers = 1
alpha = 1.0
beta = 1.0
dataaug = DataAug(256)
dataset = BacteriaDataset(
        image_folder='./Bacteria/training_data/',
        image_list_file='./Bacteria/training_data/list.txt',
        mode='train',
        transforms=dataaug,
        shuffle=shuffle
)
train_set, validation_set = torch.utils.data.random_split(dataset, [266,88])
train_loader = DataLoader(dataset=train_set, shuffle=RandomSampler(train_set), batch_size=batch_size,num_workers=num_workers,pin_memory=pin_memory)
validation_loader = DataLoader(dataset=validation_set, shuffle=shuffle, batch_size=batch_size,num_workers=num_workers, pin_memory=pin_memory)

model = UNet().to(device)
dice_loss = DiceLoss()
iou_loss = IOULoss()
# criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

def check_accuracy(loader, model):
    if loader == train_loader:
        print("Checking accuracy on training data")
    else:
        print("Checking accuracy on validation data")
    epoch_val_accuracy = 0
    epoch_val_loss = 0
    with torch.no_grad():
        for x, y in loader:
            x = x.to(device=device)
            y = y.to(device=device)
            scores = model(x)
            val_loss = alpha * dice_loss(scores, y) + beta * iou_loss(scores, y)
            acc = ((scores.argmax(dim=1) == y).float().mean())
            epoch_val_accuracy += acc/ len(loader)
            epoch_val_loss += val_loss/ len(loader)
        print('val_accuracy : {}, val_loss : {}'.format(epoch_val_accuracy,epoch_val_loss))  
    return 'val_accuracy : {}, val_loss : {}'.format(epoch_val_accuracy,epoch_val_loss)

def train():
    model.train()
    for epoch in range(num_epochs):
        loop = tqdm(train_loader, total = len(train_loader), leave = True)
        for imgs, labels in loop:
            imgs = imgs.to(device)
            labels = labels.to(device)
            outputs = model(imgs)
            loss = alpha * dice_loss(outputs, labels) + beta * iou_loss(outputs, labels)
            # loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loop.set_description(f"Epoch [{epoch}/{num_epochs}]")
            loop.set_postfix(loss = loss.item())
        if epoch % 2 == 0:
            loop.set_postfix(val_acc = check_accuracy(validation_loader, model))
    torch.save(model.state_dict(), "checkpoints/model.pth")

if __name__ == "__main__":
    train()